// javausb.cpp

#include <windows.h>

#include "javausb.h"
#include "usb.h"

// =======
// GLOBALS
// =======

// g_jobjectsCount holds a count of the number of jobjects in the vector.

static int g_jobjectsCount = 0;

// g_jobjectsLimit holds the maximum storage index of the vector's g_jobjects
// array before the vector grows this array.

static int g_jobjectsLimit = 0;

// g_jobjects holds the address of the vector's resizable jobjects array.

static jobject *g_jobjects = 0;

// ==========================
// HELPER FUNCTION PROTOTYPES
// ==========================

static bool getAttachedDevices (JNIEnv *env, const char *hubName);
static bool getHostControllerInfo (int index, char *name, char *rootHubName);
static char *getString (HANDLE handleHub, UCHAR ucPort, UCHAR ucStrIndex);

static bool vectorAppend (jobject);
static bool vectorCreate (void);
static void vectorDestroy (void);
static jobject vectorGet (int);
static int vectorSize (void);

// ========================
// JNI-ACCESSIBLE FUNCTIONS
// ========================

JNIEXPORT jobjectArray JNICALL Java_se_bluebrim_usb_example_JavaUSB_getAttachedDevices
                               (JNIEnv *env, jclass clazz, jstring hubName)
{
   // Attempt to create a vector to hold Function and Hub objects.

   if (!vectorCreate ())
       return (jobjectArray) 0;

   // Convert hub name argument from Java format to C/C++ format.

   jboolean isCopy;

   const char *chubName = env->GetStringUTFChars (hubName, &isCopy);

   // Attempt to retrieve all devices attached to the hub.

   jobjectArray devicesArray = 0;

   if (getAttachedDevices (env, chubName))
   {
       // Attempt to find the Device class.

       clazz = env->FindClass ("se/bluebrim/usb/example/Device");
       if (clazz != 0)
       {
           // Attempt to create an array of Device objects for all found
           // devices.

           devicesArray = env->NewObjectArray (vectorSize (), clazz, 0);
           if (devicesArray != 0)
           {
               // Populate array with Hub and Function objects.

               for (int i = 0; i < vectorSize (); i++)
                    env->SetObjectArrayElement (devicesArray, i, vectorGet (i));
           }
       }
   }

   // Release any memory containing copied characters.

   if (isCopy == JNI_TRUE)
       env->ReleaseStringUTFChars (hubName, chubName);

   // Cleanup vector memory.

   vectorDestroy ();

   return devicesArray;
}

JNIEXPORT jobject JNICALL Java_se_bluebrim_usb_example_JavaUSB_getHostControllerInfo
                          (JNIEnv *env, jclass clazz, jint index)
{
   // Attempt to obtain the host controller name and root hub name for the
   // specified host controller.

   char name [256];
   char rootHubName [256];

   if (!getHostControllerInfo (index, name, rootHubName))
       return (jobject) 0;

   // Attempt to find the HCInfo class.

   clazz = env->FindClass ("se/bluebrim/usb/example/HCInfo");
   if (clazz == 0)
       return (jobject) 0;

   // Attempt to allocate memory for a new HCInfo class object. Do not bother
   // to call the constructor (the default constructor does nothing).

   jobject obj = env->AllocObject (clazz);
   if (obj == 0)
       return (jobject) 0;

   // Attempt to find the name field.

   jfieldID fid = env->GetFieldID (clazz, "name", "Ljava/lang/String;");
   if (fid == 0)
       return obj;

   // Save the host controller name in the name field.

   env->SetObjectField (obj, fid, env->NewStringUTF (name));

   // Attempt to find the rootHubName field.

   fid = env->GetFieldID (clazz, "rootHubName", "Ljava/lang/String;");
   if (fid == 0)
       return obj;

   // Save the root hub name in the rootHubName field.

   env->SetObjectField (obj, fid, env->NewStringUTF (rootHubName));

   // Return the HCInfo object.

   return obj;
}

// ================
// HELPER FUNCTIONS
// ================

// --------------------
// USB Helper Functions
// --------------------

// =======================================================================
// Obtain a vector of all devices attached to a specific hub.
//
// Arguments:
//
// env     - JNI pointer, used to access JNI functions -- like FindClass()
//
// hubName - name of hub to interrogate
//
// Return:
//
// TRUE if successful; FALSE if a failure occurs
// =======================================================================

static bool getAttachedDevices (JNIEnv *env, const char *hubName)
{
   struct
   {
      ULONG ulConnectionIndex;
      ULONG ulActualLength;
      WCHAR wcName [256];
   }
   ConnectedHub;

   jobject obj;
   jclass clazz;
   HANDLE handleHub;
   DWORD dwBytesReturned;
   char attachedHubName [256];
   static SECURITY_ATTRIBUTES sa;
   USB_NODE_INFORMATION NodeInformation;
   USB_NODE_CONNECTION_INFORMATION ConnectionInformation;

   // Windows 2000 requires that a SECURITY_ATTRIBUTES structure be passed to
   // CreateFile(). Apart from specifying the structure's size, the other
   // members can keep their NULL and FALSE defaults, which were assigned via
   // the static keyword.

   sa.nLength = sizeof(SECURITY_ATTRIBUTES);

   // Attempt to obtain handle to hub. If hubName is invalid, return 
   // INVALID_HANDLE_VALUE.

   handleHub = CreateFile (hubName, GENERIC_WRITE, FILE_SHARE_WRITE, &sa,
                           OPEN_EXISTING, 0, 0);
   if (handleHub == INVALID_HANDLE_VALUE)
       return FALSE;

   // Obtain node information for the hub.

   if (!DeviceIoControl (handleHub, IOCTL_USB_GET_NODE_INFORMATION,
                         &NodeInformation, sizeof(NodeInformation),
                         &NodeInformation, sizeof(NodeInformation),
                         &dwBytesReturned, 0))
   {
      CloseHandle (handleHub);
      return FALSE;
   }

   // Enumerate all of the hub's ports.

   BOOL bStatus = TRUE;

   for (UCHAR ucPort = 1;
        ucPort <= NodeInformation.u.HubInformation.HubDescriptor.bNumberOfPorts;
        ucPort++)
   {
        // Attempt to obtain connection information for port. Ignore port if a
        // problem occurs.

        ConnectionInformation.ConnectionIndex = ucPort;
        if (!DeviceIoControl (handleHub,
                              IOCTL_USB_GET_NODE_CONNECTION_INFORMATION,
                              &ConnectionInformation,
                              sizeof(ConnectionInformation),
                              &ConnectionInformation,
                              sizeof(ConnectionInformation),
                              &dwBytesReturned, 0))
            continue;

        // Skip port if no device is connected.

        if (ConnectionInformation.ConnectionStatus != DeviceConnected)
            continue;

        // Create appropriate Device object and store appropriate information.

        if (ConnectionInformation.DeviceIsHub)
        {
            // Attempt to find Hub class.

            clazz = env->FindClass ("se/bluebrim/usb/example/Hub");
            if (clazz == 0)
            {
                bStatus = FALSE;
                break;
            }

            // Attempt to allocate memory for a new Hub object. Do not bother
            // to call the constructor (the default constructor does nothing).

            obj = env->AllocObject (clazz);
            if (obj == 0)
            {
                bStatus = FALSE;
                break;
            }

            // Attempt to find the hubName field.

            jfieldID fid = env->GetFieldID (clazz, "hubName",
                                            "Ljava/lang/String;");
            if (fid == 0)
            {
                bStatus = FALSE;
                break;
            }

            // Attempt to get the hub's name.

            ConnectedHub.ulConnectionIndex = ucPort;
            if (!DeviceIoControl (handleHub,
                                  IOCTL_USB_GET_NODE_CONNECTION_NAME,
                                  &ConnectedHub,
                                  sizeof(ConnectedHub),
                                  &ConnectedHub,
                                  sizeof(ConnectedHub),
                                  &dwBytesReturned, 0))
            {
                bStatus = FALSE;
                break;
            }

            // Convert from Unicode characters to ANSI 8-bit characters.

            lstrcpy (attachedHubName, "\\\\.\\");
            WideCharToMultiByte (CP_ACP, 0, &ConnectedHub.wcName [0],
                                 ConnectedHub.ulActualLength,
                                 &attachedHubName [4], 251, 0, 0);

            // Save the hub name string in the hubName field.

            env->SetObjectField (obj, fid, env->NewStringUTF (attachedHubName));
        }
        else
        {
            // Attempt to find Function class.

            clazz = env->FindClass ("se/bluebrim/usb/example/Function");
            if (clazz == 0)
            {
                bStatus = FALSE;
                break;
            }

            // Attempt to allocate memory for a new Function object. Do not
            // bother to call the constructor (the default constructor does
            // nothing).

            obj = env->AllocObject (clazz);
            if (obj == 0)
            {
                bStatus = FALSE;
                break;
            }
        }

        // Attempt to find the manufacturer field.

        jfieldID fid = env->GetFieldID (clazz, "manufacturer",
                                        "Ljava/lang/String;");
        if (fid == 0)
        {
            bStatus = FALSE;
            break;
        }

        // Obtain the manufacturer string.

        UCHAR ucIndex =
          ConnectionInformation.DeviceDescriptor.iManufacturer;
        char *str = getString (handleHub, ucPort, ucIndex);

        // Save the manufacturer string in the manufacturer field.

        env->SetObjectField (obj, fid, env->NewStringUTF (str));

        // Attempt to find the product field.

        fid = env->GetFieldID (clazz, "product", "Ljava/lang/String;");
        if (fid == 0)
        {
            bStatus = FALSE;
            break;
        }

        // Obtain the product string.

        ucIndex = ConnectionInformation.DeviceDescriptor.iProduct;
        str = getString (handleHub, ucPort, ucIndex);

        // Save the product string in the product field.

        env->SetObjectField (obj, fid, env->NewStringUTF (str));

        // Attempt to find the serialNumber field.

        fid = env->GetFieldID (clazz, "serialNumber",
                              "Ljava/lang/String;");
        if (fid == 0)
        {
            bStatus = FALSE;
            break;
        }

        // Obtain the serial number string.

        ucIndex = ConnectionInformation.DeviceDescriptor.iSerialNumber;
        str = getString (handleHub, ucPort, ucIndex);

        // Save the serial number string in the serialNumber field.

        env->SetObjectField (obj, fid, env->NewStringUTF (str));

        // Save Hub or Function object.

        vectorAppend (obj);
   }

   CloseHandle (handleHub);

   return bStatus;
}

// =======================================================================
// Obtain host controller information.
//
// Arguments:
//
// index       - zero-based integer index of host controller
//
// name        - buffer for host controller's name -- must be at least 255
//               characters long
//
// rootHubName - buffer for host controller's integrated root hub name --
//               must be at least 255 characters long
//
// Return:
//
// TRUE if successful; FALSE if a failure occurs
// =======================================================================

static bool getHostControllerInfo (int index, char *name, char *rootHubName)
{
   struct
   {
      ULONG ulLength;
      WCHAR wcName [256];
   }
   UnicodeName;

   HANDLE handleHC;
   char hcName [256];
   DWORD dwBytesReturned;
   static SECURITY_ATTRIBUTES sa;
   USB_NODE_INFORMATION NodeInformation;

   // Windows 2000 requires that a SECURITY_ATTRIBUTES structure be passed to
   // CreateFile(). Apart from specifying the structure's size, the other
   // members can keep their NULL and FALSE defaults, which were assigned via
   // the static keyword.

   sa.nLength = sizeof(SECURITY_ATTRIBUTES);

   // Windows 98 SE, Windows ME, and Windows 2000 use syntax \\.\HCDx (x is a 
   // 0-based integer) to name a host controller device.

   wsprintf (hcName, "\\\\.\\HCD%d", index);

   // Attempt to obtain handle to host controller. If host controller doesn't
   // exist, return INVALID_HANDLE_VALUE.

   handleHC = CreateFile (hcName, GENERIC_WRITE, FILE_SHARE_WRITE, &sa,
                          OPEN_EXISTING, 0, 0);
   if (handleHC == INVALID_HANDLE_VALUE)
       return FALSE;

   // Attempt to obtain the host controller's name.

   if (!DeviceIoControl (handleHC, IOCTL_GET_HCD_DRIVERKEY_NAME,
                         &UnicodeName, sizeof(UnicodeName),
                         &UnicodeName, sizeof(UnicodeName), &dwBytesReturned,
                         0))
   {
       CloseHandle (handleHC);
       return FALSE;
   }

   // Convert from Unicode characters to ANSI 8-bit characters.

   WideCharToMultiByte (CP_ACP, 0, &UnicodeName.wcName [0],
                        UnicodeName.ulLength, name, 255, 0, 0);

   // Attempt to obtain the root hub's name.

   if (!DeviceIoControl (handleHC, IOCTL_USB_GET_ROOT_HUB_NAME,
                         &UnicodeName, sizeof(UnicodeName),
                         &UnicodeName, sizeof(UnicodeName), &dwBytesReturned,
                         0))
   {
       CloseHandle (handleHC);
       return FALSE;
   }

   // Convert from Unicode characters to ANSI 8-bit characters.

   lstrcpy (hcName, "\\\\.\\");
   WideCharToMultiByte (CP_ACP, 0, &UnicodeName.wcName [0],
                        UnicodeName.ulLength, &hcName [4], 251, 0, 0);
   lstrcpy (rootHubName, hcName);

   CloseHandle (handleHC);

   return TRUE;
}

// ==========================================================================
// Obtain the ANSI equivalent of a Unicode string stored on a function via a
// string descriptor.
//
// Arguments:
//
// handleHub  - handle of hub to which function is attached
//
// ucPort     - index of hub port to which function is attached
//
// ucStrIndex - index of Unicode string in function's string descriptor array
//
// Return:
//
// ASCIZ string if successful; 0 if a failure occurs
// ==========================================================================

static char *getString (HANDLE handleHub, UCHAR ucPort, UCHAR ucStrIndex)
{
   bool bSuccess;
   USHORT usLangID;
   DWORD dwBytesReturned;
   USB_DESCRIPTOR_REQUEST packet;
   static unsigned char buffer [4000];

   // Attempt to get the language ID.

   memset (&packet, 0, sizeof(packet));
   packet.ConnectionIndex = ucPort;
   packet.SetupPacket.bmRequest = 0x80;
   packet.SetupPacket.bRequest = USB_REQUEST_GET_DESCRIPTOR;
   packet.SetupPacket.wValue = USB_STRING_DESCRIPTOR_TYPE << 8;
   packet.SetupPacket.wLength = 4;

   if (!DeviceIoControl (handleHub,
                         IOCTL_USB_GET_DESCRIPTOR_FROM_NODE_CONNECTION,
                         &packet, sizeof(packet), &packet,
                         sizeof(packet), &dwBytesReturned, 0))
       return 0;

   usLangID = packet.Data [2] + (packet.Data [3] << 8);

   // Attempt to get the string descriptor based on the language ID.

   memset (&packet, 0, sizeof(packet));
   packet.ConnectionIndex = ucPort;
   packet.SetupPacket.bmRequest = 0x80;
   packet.SetupPacket.bRequest = USB_REQUEST_GET_DESCRIPTOR;
   packet.SetupPacket.wValue = USB_STRING_DESCRIPTOR_TYPE << 8 | ucStrIndex;
   packet.SetupPacket.wIndex = usLangID;
   packet.SetupPacket.wLength = 255;

   if (!DeviceIoControl (handleHub,
                         IOCTL_USB_GET_DESCRIPTOR_FROM_NODE_CONNECTION,
                         &packet, sizeof(packet), &packet, sizeof(packet),
                         &dwBytesReturned, 0))
       return 0;

   // Convert from Unicode characters to ANSI 8-bit characters.

   WideCharToMultiByte (CP_ACP, 0, (LPCWSTR) &packet.Data [2], -1, buffer, 252,
                        0, 0);

   return buffer;
}

// -----------------------
// Vector Helper Functions
// -----------------------

// =============================================
// Append JNI object to vector.
//
// Arguments:
//
// o - object to append
//
// Return:
//
// TRUE if successful; FALSE if a failure occurs
// =============================================

static bool vectorAppend (jobject o)
{
   // Attempt to append jobject to vector. Grow array as necessary.

   if (g_jobjectsCount <= g_jobjectsLimit)
       g_jobjects [g_jobjectsCount++] = o;
   else
   {
       g_jobjectsLimit += g_jobjectsLimit; // Double the limit.

       jobject *jobjects =
         (jobject *) malloc (g_jobjectsLimit * sizeof(jobject *));
       if (jobjects == 0)
           return FALSE;

       for (int i = 0; i < g_jobjectsCount; i++)
            jobjects [i] = g_jobjects [i];

       free (jobjects);
       g_jobjects = jobjects;
   }

   return TRUE;
}

// =============================================
// Create the vector.
//
// Arguments:
//
// none
//
// Return:
//
// TRUE if successful; FALSE if a failure occurs
// =============================================

static bool vectorCreate (void)
{
   // Create vector with initial room for 100 jobjects before growing.

   g_jobjects = (jobject *) malloc (100 * sizeof(jobject *));

   g_jobjectsCount = 0;
   g_jobjectsLimit = 99;

   return (g_jobjects != 0) ? TRUE : FALSE;
}

// ===================
// Destroy the vector.
//
// Arguments:
//
// none
//
// Return:
//
// none
// ===================

static void vectorDestroy (void)
{
   // Release memory allocated to vector's jobjects array.

   free (g_jobjects);
}

// ========================================================================
// Return the JNI object located at the index vector position. WARNING: The
// position must be valid; otherwise a system failure will probably occur.
//
// Arguments:
//
// index - zero-based object position
//
// Return:
//
// JNI object
// ========================================================================

static jobject vectorGet (int index)
{
   // Return the jobject located in the vector at position index.

   return g_jobjects [index];
}

// =========================
// Return the vector's size.
//
// Arguments:
//
// none
//
// Return:
//
// vector size
// =========================

static int vectorSize (void)
{
   // Return number of jobjects currently in vector.

   return g_jobjectsCount;
}
